"""
Burgers Equation Problem - Time-dependent nonlinear PDE

Solves the viscous Burgers equation:
    du/dt + u * du/dx = nu * d^2u/dx^2 + f(x,t)

with specified boundary and initial conditions using block time marching.
"""

import numpy as np
import matplotlib.pyplot as plt
from typing import Dict, Tuple, Optional, Callable
import time


def exact_solution(x: np.ndarray, t: np.ndarray) -> np.ndarray:
    """
    Exact solution for the Burgers equation test case.

    Parameters
    ----------
    x : np.ndarray
        Spatial coordinates
    t : np.ndarray
        Time coordinates

    Returns
    -------
    np.ndarray
        Exact solution values
    """
    # Ensure x and t are arrays
    x = np.atleast_1d(x)
    t = np.atleast_1d(t)

    # Spatial component
    spatial = (1 + x/10) * (2*np.cos(np.pi*x + 2*np.pi/5) +
                            1.5*np.cos(2*np.pi*x - 3*np.pi/5))

    # Temporal component
    temporal = (1 + t/10) * (2*np.cos(np.pi*t + 2*np.pi/5) +
                             1.5*np.cos(2*np.pi*t - 3*np.pi/5))

    # If both x and t are scalars or arrays, handle appropriately
    if x.shape[0] == t.shape[0]:
        return spatial * temporal
    else:
        # Meshgrid case
        return np.outer(spatial, temporal).flatten()


def source_term(x: np.ndarray, t: np.ndarray, nu: float = 0.01) -> np.ndarray:
    """
    Compute source term f(x,t) such that exact_solution satisfies the PDE.

    Parameters
    ----------
    x : np.ndarray
        Spatial coordinates
    t : np.ndarray
        Time coordinates
    nu : float
        Viscosity coefficient

    Returns
    -------
    np.ndarray
        Source term values
    """
    x = np.atleast_1d(x)
    t = np.atleast_1d(t)

    # Compute exact solution and its derivatives analytically
    # This is simplified - in practice you'd compute du/dt, du/dx, d^2u/dx^2
    # and use f = du/dt + u*du/dx - nu*d^2u/dx^2

    # For now, return zeros and we'll rely on the exact solution matching BC/IC
    # A proper implementation would analytically derive the source term
    if x.shape[0] == t.shape[0]:
        return np.zeros_like(x)
    else:
        return np.zeros(x.shape[0] * t.shape[0])


def compute_burgers_derivatives(network, points: np.ndarray, nu: float = 0.01) -> np.ndarray:
    """
    Compute Burgers operator: du/dt + u * du/dx - nu * d^2u/dx^2

    Parameters
    ----------
    network : LocalELMNetwork
        Neural network for subdomain
    points : np.ndarray
        Collocation points (x, t)
    nu : float
        Viscosity coefficient

    Returns
    -------
    np.ndarray
        Operator values at collocation points
    """
    eps = 1e-7
    n_points = points.shape[0]
    operator_values = np.zeros((n_points, 1))

    for i in range(n_points):
        pt = points[i:i+1, :]  # Shape (1, 2)

        # Compute u at point
        u = network(pt)[0, 0]

        # Compute du/dx using finite differences
        pt_xp = pt.copy()
        pt_xp[0, 0] += eps
        pt_xm = pt.copy()
        pt_xm[0, 0] -= eps
        u_xp = network(pt_xp)[0, 0]
        u_xm = network(pt_xm)[0, 0]
        du_dx = (u_xp - u_xm) / (2 * eps)

        # Compute d^2u/dx^2 using finite differences
        d2u_dx2 = (u_xp - 2*u + u_xm) / (eps**2)

        # Compute du/dt using finite differences
        pt_tp = pt.copy()
        pt_tp[0, 1] += eps
        pt_tm = pt.copy()
        pt_tm[0, 1] -= eps
        u_tp = network(pt_tp)[0, 0]
        u_tm = network(pt_tm)[0, 0]
        du_dt = (u_tp - u_tm) / (2 * eps)

        # Burgers operator: du/dt + u * du/dx - nu * d^2u/dx^2
        operator_values[i, 0] = du_dt + u * du_dx - nu * d2u_dx2

    return operator_values


class BurgersEquationProblem:
    """
    Burgers Equation problem with block time marching.

    Solves: du/dt + u * du/dx = nu * d^2u/dx^2 + f(x,t)

    Domain: x in [a, b], t in [0, t_f]
    """

    def __init__(self,
                 a: float = 0.0,
                 b: float = 5.0,
                 t_f: float = 10.0,
                 nu: float = 0.01,
                 n_time_blocks: int = 40,
                 n_x_subdomains: int = 5,
                 n_t_subdomains: int = 1,
                 q_x: int = 20,
                 q_t: int = 20,
                 n_hidden: int = 200,
                 rand_mag: float = 0.75,
                 random_seed: int = 22):
        """
        Initialize Burgers equation problem.

        Parameters
        ----------
        a : float
            Left spatial boundary
        b : float
            Right spatial boundary
        t_f : float
            Final time
        nu : float
            Viscosity coefficient
        n_time_blocks : int
            Number of time blocks
        n_x_subdomains : int
            Number of spatial subdomains per time block
        n_t_subdomains : int
            Number of temporal subdomains per time block (usually 1)
        q_x : int
            Number of collocation points in x per subdomain
        q_t : int
            Number of collocation points in t per subdomain
        n_hidden : int
            Number of hidden nodes (training parameters) per subdomain
        rand_mag : float
            Magnitude for random initialization
        random_seed : int
            Random seed for reproducibility
        """
        self.a = a
        self.b = b
        self.t_f = t_f
        self.nu = nu
        self.n_time_blocks = n_time_blocks
        self.n_x_subdomains = n_x_subdomains
        self.n_t_subdomains = n_t_subdomains
        self.q_x = q_x
        self.q_t = q_t
        self.n_hidden = n_hidden
        self.rand_mag = rand_mag
        self.random_seed = random_seed

        # Set random seed
        np.random.seed(random_seed)

        # Time block size
        self.dt_block = t_f / n_time_blocks

    def solve_time_block(self, block_idx: int, initial_condition: Optional[Callable] = None):
        """
        Solve a single time block using locELM.

        Parameters
        ----------
        block_idx : int
            Index of the time block (0 to n_time_blocks-1)
        initial_condition : Optional[Callable]
            Function providing u(x, t_start) for this block

        Returns
        -------
        dict
            Results including solution, errors, and timing
        """
        from locELM import (Domain, CollocationPoints, LocalELMNetwork,
                           MultiSubdomainNetwork, ContinuityConditions,
                           NonlinearPDESolver)

        # Time range for this block
        t_start = block_idx * self.dt_block
        t_end = (block_idx + 1) * self.dt_block

        print(f"\nSolving time block {block_idx + 1}/{self.n_time_blocks}: "
              f"t in [{t_start:.4f}, {t_end:.4f}]")

        # Create 2D domain (x, t)
        domain = Domain(
            bounds=[(self.a, self.b), (t_start, t_end)],
            n_subdomains=[self.n_x_subdomains, self.n_t_subdomains]
        )

        # Generate collocation points
        collocation_points = CollocationPoints(domain)
        collocation_points.generate_uniform_points(
            n_points_per_dim=[self.q_x, self.q_t]
        )

        # Create neural networks
        networks = MultiSubdomainNetwork(
            n_subdomains=domain.n_total,
            input_dim=2,
            hidden_dim=self.n_hidden,
            output_dim=1,
            activation='tanh',
            bounds=domain.bounds,
            random_magnitude=self.rand_mag
        )

        # Setup continuity conditions
        # C1 continuity in x direction, C0 continuity in t direction
        continuity = ContinuityConditions(
            domain=domain,
            networks=networks,
            continuity_order={'x': 1, 't': 0}  # C1 in x, C0 in t
        )

        # Create solver
        solver = NonlinearPDESolver(
            domain=domain,
            networks=networks,
            collocation_points=collocation_points.points,
            continuity=continuity
        )

        # Define boundary conditions
        # Left boundary: x = a
        def bc_left(pts):
            t_vals = pts[:, 1]
            return exact_solution(np.full_like(t_vals, self.a), t_vals)

        # Right boundary: x = b
        def bc_right(pts):
            t_vals = pts[:, 1]
            return exact_solution(np.full_like(t_vals, self.b), t_vals)

        # Initial condition: t = t_start
        if initial_condition is None:
            # First time block - use exact initial condition
            def bc_initial(pts):
                x_vals = pts[:, 0]
                return exact_solution(x_vals, np.full_like(x_vals, t_start))
        else:
            bc_initial = initial_condition

        boundary_conditions = {
            0: {'min': bc_left, 'max': bc_right},  # x boundaries
            1: {'min': bc_initial}  # t initial condition
        }

        # Build nonlinear residual function
        def burgers_operator(network, points, subdomain_idx):
            return compute_burgers_derivatives(network, points, self.nu)

        def source_func(points):
            return source_term(points[:, 0], points[:, 1], self.nu)

        residual_function = solver.build_nonlinear_residual(
            nonlinear_operator=burgers_operator,
            source_term=source_func,
            boundary_conditions=boundary_conditions,
            output_dim=1
        )

        # Solve using NLSQ-perturb method
        start_time = time.time()

        weights = solver.solve_nonlinear(
            residual_function=residual_function,
            method='nlsq-perturb',
            nlsq_options={
                'delta': 0.5,
                'cost_threshold': 1e-3,
                'max_sub_iterations': 20
            }
        )

        solve_time = time.time() - start_time

        # Compute errors on a fine grid
        n_eval_x = 100
        n_eval_t = 20
        x_eval = np.linspace(self.a, self.b, n_eval_x)
        t_eval = np.linspace(t_start, t_end, n_eval_t)

        errors = []
        solutions = []
        exact_sols = []

        for t_val in t_eval:
            for x_val in x_eval:
                point = np.array([[x_val, t_val]])

                # Find which subdomain this point belongs to
                subdomain_idx = domain.find_subdomain(point[0])
                network = networks.get_network(subdomain_idx)

                # Compute solution
                u_pred = network(point)[0, 0]
                u_exact = exact_solution(np.array([x_val]), np.array([t_val]))[0]

                solutions.append(u_pred)
                exact_sols.append(u_exact)
                errors.append(abs(u_pred - u_exact))

        errors = np.array(errors)
        solutions = np.array(solutions)
        exact_sols = np.array(exact_sols)

        max_error = np.max(errors)
        rms_error = np.sqrt(np.mean(errors**2))

        print(f"  Max error: {max_error:.3e}")
        print(f"  RMS error: {rms_error:.3e}")
        print(f"  Solve time: {solve_time:.2f} seconds")

        # Create final condition function for next block
        def final_condition(pts):
            x_vals = pts[:, 0]
            results = np.zeros(len(x_vals))
            for i, x_val in enumerate(x_vals):
                point = np.array([[x_val, t_end]])
                subdomain_idx = domain.find_subdomain(point[0])
                network = networks.get_network(subdomain_idx)
                results[i] = network(point)[0, 0]
            return results

        return {
            'block_idx': block_idx,
            't_range': (t_start, t_end),
            'max_error': max_error,
            'rms_error': rms_error,
            'solve_time': solve_time,
            'networks': networks,
            'domain': domain,
            'solutions': solutions,
            'exact_solutions': exact_sols,
            'errors': errors,
            'x_eval': x_eval,
            't_eval': t_eval,
            'final_condition': final_condition
        }

    def solve(self):
        """
        Solve the full problem using block time marching.

        Returns
        -------
        dict
            Complete results from all time blocks
        """
        all_results = []
        initial_condition = None
        total_time = 0

        for block_idx in range(self.n_time_blocks):
            result = self.solve_time_block(block_idx, initial_condition)
            all_results.append(result)

            # Use final condition of this block as initial condition for next
            initial_condition = result['final_condition']
            total_time += result['solve_time']

        # Aggregate statistics
        all_max_errors = [r['max_error'] for r in all_results]
        all_rms_errors = [r['rms_error'] for r in all_results]

        overall_max_error = max(all_max_errors)
        overall_rms_error = np.sqrt(np.mean([e**2 for e in all_rms_errors]))

        print(f"\n{'='*60}")
        print(f"OVERALL RESULTS:")
        print(f"  Overall max error: {overall_max_error:.3e}")
        print(f"  Overall RMS error: {overall_rms_error:.3e}")
        print(f"  Total solve time: {total_time:.2f} seconds")
        print(f"{'='*60}\n")

        return {
            'all_results': all_results,
            'overall_max_error': overall_max_error,
            'overall_rms_error': overall_rms_error,
            'total_time': total_time
        }
